# -*- coding: utf-8 -*-
# SimpleItemCF.py 基于物品

from AnimeListLoader import AnimeListLoader
from surprise import KNNBasic
import heapq
from collections import defaultdict
from operator import itemgetter

testSubject = '85'
k = 10

ml = AnimeListLoader()
data = ml.load_dataset()

trainSet = data.build_full_trainset()

sim_options = {
    'name': 'cosine',
    'user_based': False
}

# 模型是KNN
model = KNNBasic(sim_options=sim_options)
# 训练模型, 拟合数据
model.fit(trainSet)
# 计算相似度矩阵
simsMatrix = model.compute_similarities()

# 将用户原始ID（字符串类型）转换为算法内部使用的连续整数ID
testUserInnerID = trainSet.to_inner_uid(testSubject)

# 得到85号的评分数据
testUserRatings = trainSet.ur[testUserInnerID]
# 找到k个评分最高的
kNeighbors = heapq.nlargest(k, testUserRatings,
                            # t[0]是算法内部的整数ID, t[1]则是用户对物品的评分
                            key=lambda t: t[1])

# 获取与我们喜欢的项目相似的项目（按评分加权）
candidates = defaultdict(float)
for itemID, rating in kNeighbors:
    similarityRow = simsMatrix[itemID]  # 找到和被评分物品相似的
    for innerID, score in enumerate(similarityRow):
        # 记录新物品的评分
        candidates[innerID] += score * (rating / 5.0)

# 构建一个字典，记录该用户已经看过的项目
watched = {}  # 记录85号看过的
for itemID, rating in trainSet.ur[testUserInnerID]:
    watched[itemID] = 1

# 获取来自相似项目的最高评分项目：
pos = 0
for itemID, ratingSum in sorted(candidates.items(), key=itemgetter(1), reverse=True):
    if not itemID in watched:
        movieID = trainSet.to_raw_iid(itemID)
        print(ml.get_name(int(movieID)), ratingSum)
        pos += 1
        if (pos > 10):
            break
